package com.yupi.yuaiagent.agent;

import com.yupi.yuaiagent.agent.model.AgentState;
import io.micrometer.common.util.StringUtils;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;

/**
 * 抽象基础代理类，用来管理代理状态和执行流程
 *
 * 提供状态转换，内存管理和基于步骤的执行循环的基础功能
 *
 * 子类必须实现step（）方法，用来执行具体的代理逻辑
 */

@Slf4j
@Data
public abstract class BaseAgent {

    //核心属性
    private String name;

    //提示
    private String systemPrompt;
    private String nextStepPrompt;

    //状态
    private AgentState state =AgentState.IDLE;

    //执行空值
    private int maxSteps = 10;
    private int currentStep = 0;

    //消息列表（自主维护会话上下文）
    private String[] messages;

    private ChatClient chatClient;

    //执行逻辑
    public List<Message> messageList = new ArrayList<>();

    /**
     * 执行代理逻辑
     * @param userPrompt 用户提示词
     * @return 执行结果
     */
    public String run(String userPrompt){
        if (this.state !=AgentState.IDLE) {
            throw new RuntimeException("Cannot run agent from state: " + this.state);
        }

        if(StringUtils.isBlank(userPrompt)) {
            throw new RuntimeException("Cannot run agent with empty user prompt");
        }

        //更改状态为运行中
        this.state = AgentState.RUNNING;

        //记录上下文(含用户的输入信息)
        messageList.add(new UserMessage(userPrompt));

        //保存结果列表
        List<String> results = new ArrayList<>();

        //执行循环
        try {
            for (int i = 0; i < maxSteps && state!=AgentState.FINISHED; i++) {
                int stepNumber = i + 1;
                currentStep = stepNumber;
                log.info("当前执行的步骤currentStep {} of {}", stepNumber, maxSteps);

                // 执行步骤
                String stepResult = step();
                // 记录结果
                String result = "Step: " + currentStep + ":" + stepResult;
                results.add(result);
            }
            if (currentStep >= maxSteps) {
                state = AgentState.FINISHED;
                results.add("Terminated : Reached max steps (" + maxSteps + ")");
            }
            return String.join("\n", results);
        } catch (Exception e) {
            state = AgentState.ERROR;
            log.error("Error running agent", e);
            return "Error running agent: " + e.getMessage();
        } finally {
            // 释放资源
            this.cleanUp();
        }
    }
    /**
     * 运行代理（流式输出）
     * @param userPrompt 用户提示词
     * @return 执行结果
     */
    public SseEmitter runStream(String userPrompt) {
        // 创建一个超时时间为5分钟的SseEmitter
        SseEmitter sseEmitter = new SseEmitter(5 * 60 * 1000L);
        // 使用线程异步处理，避免阻塞主线程
        CompletableFuture.runAsync(() -> {
            // 1. 基础校验
            try {
                if (this.state != AgentState.IDLE) {
                    sseEmitter.send("Cannot run agent from state: " + this.state);
                    sseEmitter.complete();
                    return;
                }
                if (StringUtils.isBlank(userPrompt)) {
                    sseEmitter.send("Cannot run agent with empty user prompt");
                    sseEmitter.complete();
                    return;
                }
            } catch (Exception e) {
                sseEmitter.completeWithError(e);
            }
            // 2. 执行代理逻辑，更新状态
            this.state = AgentState.RUNNING;
            // 3. 记录上下文（含用户的输入信息）
            messageList.add(new UserMessage(userPrompt));
            ArrayList<Object> results = new ArrayList<>();
            try {
                // 执行循环
                for (int i = 0; i < maxSteps && state != AgentState.FINISHED; i++) {
                    int stepNumber = i + 1;
                    currentStep = stepNumber;
                    log.info("Executing step {} of {}", stepNumber, maxSteps);
                    // 执行步骤
                    String stepResult = step();
                    // 记录结果
                    String result = "Step: " + currentStep + ":" + stepResult;
                    results.add(result);
                    // 发送sse消息
                    sseEmitter.send(result);
                }
                if (currentStep >= maxSteps) {
                    state = AgentState.FINISHED;
                    results.add("Terminated : Reached max steps (" + maxSteps + ")");
                    sseEmitter.send("Terminated : Reached max steps (" + maxSteps + ")");
                }
                sseEmitter.complete();
            } catch (Exception e) {
                state = AgentState.ERROR;
                log.error("Error running agent", e);
                try {
                    sseEmitter.send("Error running agent: " + e.getMessage());
                    sseEmitter.complete();
                } catch (IOException ex) {
                    sseEmitter.completeWithError(ex);
                }
            } finally {
                // 释放资源
                this.cleanUp();
            }
        });
        // 设置超时回调
        sseEmitter.onTimeout(() -> {
            this.state = AgentState.ERROR;
            this.cleanUp();
            log.warn("Timeout running agent");
        });

        // 返回SseEmitter
        sseEmitter.onCompletion(() -> {
            if (this.state == AgentState.RUNNING) {
                this.state = AgentState.FINISHED;
            }
            this.cleanUp();
            log.info("SSE connection completed");
        });
        return sseEmitter;
    }

    protected void cleanUp() {
        // 可以让子类来重写此方法来释放资源
    }

    /**
     * 定义单个步骤的逻辑
     * @return 执行结果
     */
    public abstract String step();
}
